from http.server import HTTPServer, SimpleHTTPRequestHandler
import os
import xml.etree.ElementTree as ET
import subprocess
import json
import tempfile
from main import node_to_dict, dict_to_spans
import time

class CustomHandler(SimpleHTTPRequestHandler):
    def do_GET(self):
        if self.path == '/':
            self.handle_root()
        elif self.path == '/screen.png':
            self.handle_image()
        elif self.path == '/refresh':
            self.handle_refresh()
        else:
            super().do_GET()

    def handle_root(self):
        try:
            start_time = time.time()
            
            # 获取节点信息
            nodes_dict = self.get_nodes()
            
            # 读取模板
            with open('template.html', 'r', encoding='utf-8') as f:
                template = f.read()
            
            # 生成树内容
            tree_content = dict_to_spans(nodes_dict)
            
            # 替换模板内容
            html_content = template.replace('{tree_content}', tree_content)
            
            # 计算耗时
            elapsed_time = round((time.time() - start_time) * 1000)  # 转换为毫秒
            
            # 发送响应
            self.send_response(200)
            self.send_header('Content-type', 'text/html; charset=utf-8')
            self.send_header('X-Render-Time', str(elapsed_time))  # 添加耗时头
            self.end_headers()
            self.wfile.write(html_content.encode('utf-8'))
            
        except Exception as e:
            self.send_error(500, str(e))

    def handle_image(self):
        try:
            # 获取新截图
            subprocess.run('adb shell screencap -p /sdcard/screen.png', shell=True, check=True)
            subprocess.run('adb pull /sdcard/screen.png screen.png', shell=True, check=True)
            
            # 读取并发送图片
            with open('screen.png', 'rb') as f:
                image_data = f.read()
            
            self.send_response(200)
            self.send_header('Content-type', 'image/png')
            self.end_headers()
            self.wfile.write(image_data)
            
            # 清理临时文件
            try:
                os.remove('screen.png')
            except:
                pass
            
        except Exception as e:
            self.send_error(500, str(e))

    def handle_refresh(self):
        try:
            # 获取新的节点信息和截图
            nodes_dict = self.get_nodes()
            tree_content = dict_to_spans(nodes_dict)
            
            # 返回成功响应
            self.send_response(200)
            self.send_header('Content-type', 'application/json')
            self.send_header('Access-Control-Allow-Origin', '*')
            self.end_headers()
            self.wfile.write(json.dumps({
                'status': 'success',
                'tree_content': tree_content
            }).encode())
            
        except Exception as e:
            self.send_error(500, str(e))

    def get_nodes(self):
        """获取节点信息"""
        subprocess.run('adb shell uiautomator dump /sdcard/window_dump.xml', shell=True, check=True)
        
        temp_dir = tempfile.gettempdir()
        temp_xml = os.path.join(temp_dir, 'window_dump.xml')
        
        subprocess.run(f'adb pull /sdcard/window_dump.xml "{temp_xml}"', shell=True, check=True)
        
        tree = ET.parse(temp_xml)
        root = tree.getroot()
        nodes_dict = node_to_dict(root)
        
        try:
            os.remove(temp_xml)
        except:
            pass
            
        return nodes_dict

def run_server(port=8000):
    server_address = ('', port)
    httpd = HTTPServer(server_address, CustomHandler)
    print(f"服务器启动在 http://localhost:{port}")
    httpd.serve_forever()

if __name__ == '__main__':
    run_server() 